import pandas as pd #importing all the important packages
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
plt.style.use('fivethirtyeight')

import streamlit as st

def pre_data():
    df = pd.read_csv('data/Pokemon.csv')
    df.columns = df.columns.str.upper().str.replace('_', '')
    df[df['LEGENDARY']==True].head(5) #Showing the legendary pokemons
    df = df.set_index('NAME') 
    df.index = df.index.str.replace(".*(?=Mega)", "")
    df=df.drop(['#'],axis=1) 
    df['TYPE 2'].fillna(df['TYPE 1'], inplace=True)
    return df

df = pre_data()

def plot_scatter():
    pk_type1  = st.selectbox("选择你要对比的第一种类型的宝可梦",df['TYPE 1'].unique())
    pk_type2  = st.selectbox("选择你要对比的第二种类型的宝可梦",df['TYPE 1'].unique())
    
    fire=df[(df['TYPE 1']==pk_type1) | ((df['TYPE 2'])==pk_type1)] #fire contains all fire pokemons
    water=df[(df['TYPE 1']==pk_type2) | ((df['TYPE 2'])==pk_type2)] #all water pokemins
    plt.scatter(fire.ATTACK.head(50),fire.DEFENSE.head(50),color='Red',label=pk_type1,marker="*",s=50) #scatter 
    plt.scatter(water.ATTACK.head(50),water.DEFENSE.head(50),color='Blue',label=pk_type2,s=25)
    plt.xlabel("Attack")
    plt.ylabel("DEFENCE")
    plt.legend()
    # plt.plot()
    fig=plt.gcf() #get the current figure using .gcf()
    fig.set_size_inches(12,6) #set the size for the figure
    # plt.show()
    
    st.pyplot(plt)
    
    